#!/usr/bin/python
# Author: Jimmy Saw
# Date: 06-21-2011
# This script extracts the paired-end reads and singletons from the original multi-fasta file.
# It uses the "trim point" coordinates to trim and keep only the good portion needed for assembly.
# Prints only sequences >= 10bp.
# 
# The program needs 2 files to run (and needs Biopython obviously):
# 1. Multi-fasta file of original 454 reads file
# 2. 454TrimStatus.txt file
# Usage: python extract454Reads.py multifasta.fasta 454TrimStatus.txt
#
# Logic:
#          ---->                                ---->
# {------454_left------------}{Linker}{--------454_right--------}
# {Actual Right end of insert}{Linker}{Actual Left end of insert}
# {Reverse complement this   }{remove}{Keep as it is            }

import sys
import re
from Bio import SeqIO

seqfile = sys.argv[1]
sf = SeqIO.parse(seqfile, "fasta")

trimfile = sys.argv[2]
tf = open(trimfile, "rU")
tfl = tf.readlines()

seqs_dict = {}

m1 = re.compile('(\w+)_left')
m2 = re.compile('(\w+)_right')
m3 = re.compile('(\w+)')

for seq in sf:
    seqs_dict[seq.id] = seq

for line in tfl[1:]: #Skip the header line
    t = line.split('\t')
    accnum = t[0]
    trimpoints_used = t[1].split('-')
    trim_start = int(trimpoints_used[0])
    trim_stop = int(trimpoints_used[1])
    outseq_length = trim_stop - trim_start + 1
    used_trim_length = t[2]
    orig_trimpoints = t[3]
    orig_trim_length = t[4]
    raw_length = t[5]
    seqname = ""
    outseq = ""
    paired = "no"
    position = ""
    outread_name = ""
    if m1.match(accnum):
        p = m1.match(accnum)
        seqname = p.group(1)
        paired = "yes"
        position = "R"
    elif m2.match(accnum):
        p = m2.match(accnum)
        seqname = p.group(1)
        paired = "yes"
        position = "L"
    elif m3.match(accnum):
        p = m3.match(accnum)
        seqname = p.group(1)
        paired = "no"
    if seqname in seqs_dict:
        if outseq_length >= 10:
            if position == "R": #need to reverse complement
                outread_name = seqname + "_R"
                tmp = seqs_dict[seqname].seq
                tmp2 = tmp[trim_start:trim_stop]
                outseq = tmp2.reverse_complement()
            elif position == "L": #keep as it is
                tmp = seqs_dict[seqname].seq
                outseq = tmp[trim_start:trim_stop]
                outread_name = seqname + "_L"
                #outseq = tmp2
            elif position == "": #keep as it is
                outread_name = seqname
                tmp = seqs_dict[seqname].seq
                outseq = tmp[trim_start:trim_stop]
            print ">" + outread_name
            print outseq

tf.close()